#!/usr/bin/env python
#Copyright (c) 2011 Isaac Dawson (WBTS Project)
#Permission is hereby granted, free of charge, to any person obtaining a copy 
#of this software and associated documentation files (the "Software"), to deal 
#in the Software without restriction, including without limitation the rights 
#to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 
#copies of the Software, and to permit persons to whom the Software is furnished
#to do so, subject to the following conditions:
#
#The above copyright notice and this permission notice shall be included in all
#copies or substantial portions of the Software.
#
#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 
#LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 
#THE SOFTWARE.
try:
    import txmongo
except ImportError, msg:
    print '-'*80
    print "Error unable to find/import txmongo! MongoDB will not work!"
    print "Please install txmongo if you wish to use MongoDB."
    print '-'*80
    pass
from twisted.internet import defer
from twisted.python import components, log

import storage
import wbts.utils.mgmt_utils

class WBTSMongoDB(storage.WBTSStorage):
    """
    We remap the databases from the txmongo connection to our own
    class.
    """
    mDbConn = None
    def __init__(self, connect_string):
        self.mDbHost, self.mDbPort, self.mDbName = self._parseConnectionString(
            connect_string)
        
        
    def connect(self):        
        d = txmongo.MongoConnection(self.mDbHost, self.mDbPort)
        d.addCallback(self._assignConnection)
        d.addErrback(self._dbError)
        return d
    
    def _assignConnection(self, conn):
        self.mDbConn = conn
        return 0

    def _dbError(self, err_msg):
        log.err(err_msg)
        return 1
    
    def disconnect(self):
        if (self.mDbConn == None):
            return defer.succeed(0)
        return self.mDbConn.disconnect()
        
    def createTable(self, tbl_name, tbl_callable=None):
        """    
    Creates a new table and assigns it as an attribute. If the table already
    exists (the attribute has been set) then it returns a deferred of 1.
    @returns
    (deferred) 0 - table created.
    (deferred) 1 - already exists
        """
        if (tbl_name == '' or not isinstance(tbl_name, str)):
            return defer.fail(
                storage.WBTSStorageException(
                    'table name is not a string or empty.')
                )
        
        if (self.mDbConn == None):
            return defer.fail(
                storage.WBTSStorageException(
                    'We do not have a valid DB connection.')
            )
        
        try:
            getattr(self, tbl_name)
            return defer.succeed(1)
        except AttributeError:
            pass
        print self.mDbConn
        setattr(self, tbl_name, WBTSMongoTable(self.mDbConn[self.mDbName],
                                               tbl_name))
        
        return defer.succeed(0)

    def dropTable(self, tbl_name):
        ref = None
        try:
            ref = getattr(self, tbl_name)
            delattr(self, tbl_name)
            return ref.drop()
        except AttributeError:
            pass
        return defer.succeed(1)
    
    def _dropTable(self, result):
        return
    
    def _parseConnectionString(self, connect_string):
        try:
            host,port,name = connect_string.split(':')
            port = int(port)
            ret = wbts.utils.mgmt_utils.is_port_valid(port)
            if (ret == True):
                return host,port,name
        except ValueError, msg:
            pass

        print "Unable to parse MongoDB connect string, it must be in the "
        print "following format: ip.ip.ip.ip:port:dbname"
        return None,None,None

class WBTSMongoTable(storage.WBTSTableObject):
    """
    Not totally necessary but there are some things about the pymongo Collection
    object we do not need. Also it removes any doubt of possible inconsistency
    between the various data storage types.
    """
    def __init__(self, mongo_db, tbl_name):
        self.db = mongo_db
        self.name = tbl_name
        return
    
    def find(self, search_key=None, skip=0, limit=-1):
        """
    Finds a splice of records (starting at skip up to skip+limit). 
    @params
    search_key - A dictionary of keys/values to be used to find records.
    skip - an integer of the # of records to skip ahead
    limit - the number of records to return, if -1 returns rest of records
            starting from skip.
    @returns
    A deferred which will return a list object of records between skip
    and skip+limit Ex: (table.records[skip:skip+limit]) or records which match
    the values from search_key.
        """
        d = None
        if (search_key != None):
            d = self.db[self.name].find(search_key)
        # in this case we just need to find *, but mongodb will only
        # return 1 record...
        elif (search_key == None and (skip == 0) and (limit == -1)):
            d = self.db[self.name].find()
        else:
            d = self.db[self.name].find(skip=skip, limit=limit)
        d.addCallback(self._find)
        d.addErrback(self._fail)
        return d
        
    def _find(self, result):
        if (isinstance(result, list)):
            if (len(result) >= 1 ):
                if result[0].has_key('$err'):
                    log.msg("find returned error: %s"%(result[0]['$err'],))
                    return None
        return result
    
    def find_one(self, search_key=None):
        """
    Find a single record
    @params
    search_key - A dictionary of keys/values to be used to find a record
    @returns
    A deferred which results in a dict object of the first record found
    or an empty dict {} if nothing found.
        """
        d = self.db[self.name].find_one(search_key)
        d.addCallback(self._find)
        d.addErrback(self._fail)
        return d
    
    
    def insert(self, record, safe=True):
        """
    Inserts a record (dictionary object) of data into the table.
    @params
    record - a dictionary object of our data. Keep in mind SQL versions will
    be strict about the data type, pydictdb and mongodb don't care :).
    @returns
    A deferred of 0 for success and 1 for failure.
        """
        d = self.db[self.name].insert(record, safe=True)
        d.addCallback(self._insert)
        d.addErrback(self._fail)
        return d
    
    def _insert(self, result):
        """
    Returns simply 0 for success and 1 for failure instead of ObjectIds which
    the original Pymongo Collection object does.
        """
        return 0
    
    def _fail(self, result):
        log.err(result)
        return 1
    
    def remove(self, search_key=None, safe=True):
        """
    Removes ANY record that maches the values from search_key.
    @params
    search_key - a dictionary of key/values to be used to find a record.
    @returns
    A deferred of 0 for success and 1 for failure or record not found.
        """
        if (search_key == None):
            search_key = {}
        
        d = self.db[self.name].remove(search_key, safe)
        d.addCallback(self._remove)
        d.addErrback(self._fail)
        return d
    
    def _remove(self, result):
        return 0
    

    def count(self):
        """
    Get the number of records
        """
        return self.db[self.name].count()
    
    def drop(self):
        """
    Drops the current database table.
    @returns
    A deferred value 0 for success 1 for error
        """
        d = self.db[self.name].drop(safe=True)
        d.addCallback(self._drop)
        d.addErrback(self._fail)
        return d
    
    def _drop(self, result):
        return 0
    
